import cv2
import numpy as np
import time
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 读取图像
img_path = "test_image/test_6/01.jpg"
img = cv2.imread(img_path)

if img is None:
    raise FileNotFoundError(f"无法加载图像：{img_path}，请检查路径或文件名！")

# 灰度图
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

# 1. Harris 角点检测（OpenCV 版）
def harris_opencv(gray):
    start = time.time()

    gray_f = np.float32(gray)
    dst = cv2.cornerHarris(src=gray_f, blockSize=2, ksize=3, k=0.05)
    dst = cv2.dilate(dst, None)

    res = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
    res[dst > 0.01 * dst.max()] = [0, 0, 255]

    elapsed = time.time() - start
    return res, elapsed

# 2. Harris 角点检测（手写版）
def harris_manual(gray):
    start = time.time()

    gray_f = np.float32(gray) / 255.0

    # 1. Sobel 计算梯度
    Ix = cv2.Sobel(gray_f, cv2.CV_32F, 1, 0, ksize=3)
    Iy = cv2.Sobel(gray_f, cv2.CV_32F, 0, 1, ksize=3)

    Ixx = Ix * Ix
    Iyy = Iy * Iy
    Ixy = Ix * Iy

    # 2. 高斯平滑
    Sxx = cv2.GaussianBlur(Ixx, (5, 5), 1)
    Syy = cv2.GaussianBlur(Iyy, (5, 5), 1)
    Sxy = cv2.GaussianBlur(Ixy, (5, 5), 1)

    k = 0.04
    detM = Sxx * Syy - Sxy * Sxy
    traceM = Sxx + Syy
    R = detM - k * (traceM ** 2)

    # 归一化到 [0, 1]，保留 float32 精度
    R_norm = cv2.normalize(R, None, 0, 1.0, cv2.NORM_MINMAX)

    # 阈值
    threshold = 0.25

    res = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)

    h, w = gray.shape

    for y in range(2, h - 2):
        for x in range(2, w - 2):

            if R_norm[y, x] <= threshold:
                continue

            window = R_norm[y-4:y+5, x-4:x+5]   # 9×9


            if window.size == 0:
                continue

            if R_norm[y, x] == window.max():
                cv2.circle(res, (x, y), 1, (0, 0, 255), -1)

    elapsed = time.time() - start
    return res, elapsed



# 3. SUSAN 角点检测（手写版）
#    经典 SUSAN：圆模板 + USAN 面积
def susan_manual(gray, r=3, t=20, g=0.75):
    """
    r:   圆形模板半径
    t:   亮度相似阈值 (|I(q) - I(p)| < t 认为相似)
    g:   比例系数，用于控制角点判定阈值
    """
    start = time.time()

    gray_f = np.float32(gray)
    h, w = gray.shape

    # 圆形模板坐标偏移
    mask_coords = []
    for dy in range(-r, r + 1):
        for dx in range(-r, r + 1):
            if dy * dy + dx * dx <= r * r:
                mask_coords.append((dy, dx))
    N = len(mask_coords)  # 模板内像素个数

    usan = np.zeros_like(gray_f, dtype=np.float32)

    # 逐像素计算 USAN 面积
    for y in range(r, h - r):
        for x in range(r, w - r):
            center = gray_f[y, x]
            similar = 0
            for dy, dx in mask_coords:
                if abs(gray_f[y + dy, x + dx] - center) < t:
                    similar += 1
            usan[y, x] = similar

    # 角点响应：USAN 越小越像角点，所以用 R = N - USAN
    R = N - usan

    # 归一化方便阈值和显示
    R_norm = cv2.normalize(R, None, 0, 255, cv2.NORM_MINMAX)
    R_norm = np.uint8(R_norm)

    res = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)

    # 经验阈值
    threshold = 0.7 * R_norm.max()

    for y in range(r + 1, h - r - 1):
        for x in range(r + 1, w - r - 1):
            if R_norm[y, x] > threshold:
                # 简单 3x3 非极大值抑制
                if R_norm[y, x] == np.max(R_norm[y-1:y+2, x-1:x+2]):
                    cv2.circle(res, (x, y), 1, (0, 255, 0), -1)

    elapsed = time.time() - start
    return res, elapsed



# 运行各算法
harris_cv_img, t_harris_cv = harris_opencv(gray)
harris_manual_img, t_harris_manual = harris_manual(gray)
susan_img, t_susan_manual = susan_manual(gray)

print("===== 处理时间对比 =====")
print(f"Harris (OpenCV)   : {t_harris_cv:.4f} s")
print(f"Harris (手写)     : {t_harris_manual:.4f} s")
print(f"SUSAN  (手写)     : {t_susan_manual:.4f} s")

# 显示结果（用 matplotlib）
plt.figure(figsize=(12, 9))

plt.subplot(2, 2, 1)
plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
plt.title("原图")
plt.axis('off')

plt.subplot(2, 2, 2)
plt.imshow(cv2.cvtColor(harris_cv_img, cv2.COLOR_BGR2RGB))
plt.title(f"Harris OpenCV (t={t_harris_cv:.3f}s)")
plt.axis('off')

plt.subplot(2, 2, 3)
plt.imshow(cv2.cvtColor(harris_manual_img, cv2.COLOR_BGR2RGB))
plt.title(f"Harris 手写 (t={t_harris_manual:.3f}s)")
plt.axis('off')

plt.subplot(2, 2, 4)
plt.imshow(cv2.cvtColor(susan_img, cv2.COLOR_BGR2RGB))
plt.title(f"SUSAN 手写 (t={t_susan_manual:.3f}s)")
plt.axis('off')

plt.tight_layout()
plt.show()
